import logging
import time, datetime
import sys
import uuid
import os

from sqlalchemy import *
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session, relationship, backref
from sqlalchemy.event import listen
from contextlib import contextmanager

from lxml import etree

import netadmin
import report
import utils

import ovas
import nm


from config import *

logger = logging.getLogger('dmzwatch.'+__name__)

engine = create_engine(DMZWATCH_DB, echo=DB_DEBUG, pool_timeout=30)
Base = declarative_base(bind=engine)
Session = scoped_session(sessionmaker(engine))

                    

class GUID(TypeDecorator):
    """Platform-independent GUID type.

    Uses Postgresql's UUID type, otherwise uses
    CHAR(32), storing as stringified hex values.

    """
    impl = CHAR

    def load_dialect_impl(self, dialect):
        if dialect.name == 'postgresql':
            return dialect.type_descriptor(UUID())
        else:
            return dialect.type_descriptor(CHAR(32))

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name == 'postgresql':
            return str(value)
        else:
            if not isinstance(value, uuid.UUID):
                return "%.32x" % uuid.UUID(value)
            else:
                # hexstring
                return "%.32x" % value

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        else:
            return uuid.UUID(value)                 
                 
                 
class Incident(Base):
    __tablename__ = 'incidents'
    incident_id = Column(GUID, primary_key=True)
    
    incident_class = Column('incident_class', String(50))
    
    type = Column('type', String(VARCHAR_LENGTH))
    severity = Column('severity', String(VARCHAR_LENGTH))
    description = Column('description', UnicodeText(TEXT_LENGTH))
    date = Column('date', DateTime, nullable=False)
    __mapper_args__ = {'polymorphic_on': incident_class}
    
    def __init__ (self, guid, type, severity, description):
        self.incident_id = guid
        self.date = datetime.datetime.now()
        self.type = type
        self.severity = severity
        self.description = description
        self._send()
    
    def   _send(self):
        report.send(self.type, self.severity, self.description)     


class Service(Incident):
    __tablename__ = 'services'
    __mapper_args__ = {'polymorphic_identity': 'service'}
    service_id = Column(GUID, ForeignKey('incidents.incident_id', ondelete='CASCADE'), primary_key=True)
    scan_id = Column(Integer, ForeignKey('scans.scan_id', ondelete='CASCADE'))

    ip = Column('ip', String(VARCHAR_LENGTH))
    port =  Column('port', Integer)
    proto = Column('proto', String(VARCHAR_LENGTH))
    reason = Column('reason', String(VARCHAR_LENGTH))
    name = Column('name', String(VARCHAR_LENGTH))
    product = Column ('product', String(VARCHAR_LENGTH))
    version = Column('version', String(VARCHAR_LENGTH))
    extrainfo = Column('extra', String(VARCHAR_LENGTH))
    state = Column('state', String(VARCHAR_LENGTH))
    script = Column('script', UnicodeText(TEXT_LENGTH))
    ext_avail = Column('ext', Boolean)

    scan = relationship("Scan", backref = "services")
    
    def __init__ (self, ip, port, proto, reason, name, product, version, extrainfo, state, script, link = 'none'):
        self.service_id = uuid.uuid1()
        self.ip = ip
        self.port = port
        self.proto = proto
        self.reason = reason
        self.name = name
        self.product = product
        self.version = version
        self.extrainfo = extrainfo
        self.state = state
        self.script = script
        
        type = IncidentType.dmzwatch_service
        severity = Severity.medium
        
        description = '{} service found by {} at {} {}:{}, {}:{}:{}; extra: {}; {}; FULL REPORT AT {}'\
                        .format(state, reason, ip, proto, port, name, product, version, extrainfo, script, link )
        logger.debug('%s', description)                
        super().__init__(self.service_id, type, severity, description)
    
    
trig_ddl = DDL("CREATE TRIGGER service_delete AFTER DELETE ON services FOR EACH ROW BEGIN DELETE FROM incidents where incident_id = old.service_id; END")
tbl = Service.__table__
listen(tbl, 'after_create', trig_ddl) 
    
class Vulnerability(Incident):
    __tablename__ = 'vulns'
    __mapper_args__ = {'polymorphic_identity': 'vulnerability'}
    vuln_id = Column(GUID, ForeignKey('incidents.incident_id', ondelete='CASCADE'), primary_key=True)
    scan_id = Column(Integer, ForeignKey('scans.scan_id', ondelete='CASCADE'))
    
    host  = Column('host', String(VARCHAR_LENGTH))
    port = Column('port', String(VARCHAR_LENGTH))
    threat = Column('threat', String(VARCHAR_LENGTH))
    name = Column('name', String(VARCHAR_LENGTH))
    family = Column('family', String(VARCHAR_LENGTH))
    cvss = Column('cvss', String(VARCHAR_LENGTH))
    risk = Column('risk', String(VARCHAR_LENGTH))
    bid = Column('bid', String(VARCHAR_LENGTH))
    cert = Column('cert', String(VARCHAR_LENGTH))
    cve = Column('cve', String(VARCHAR_LENGTH))
    tags = Column('tags', UnicodeText(TEXT_LENGTH))
    xref = Column('xref', UnicodeText(TEXT_LENGTH))
    description = Column('description', UnicodeText(TEXT_LENGTH))
    link = Column('link', String(VARCHAR_LENGTH))
 
    scan = relationship("Scan", backref = "vulns")
    
    def __init__(self, host, port, threat, name, family, cvss_base, risk_factor, bid, cert, cve, tags, xref, description, link):
        self.vuln_id = uuid.uuid1()
        self.host = host
        self.port = port
        self.threat = threat
        self.name = name
        self.family = family
        self.cvss = cvss_base
        self.risk = risk_factor
        self.bid = bid
        self.cert = cert
        self.cve = cve
        self.tags = tags
        self.xref = xref
        self.description = description
        
        type = IncidentType.dmzwatch_vuln
        severity = self.threat
        message = '{}:{} vulnerability found at {}:{}, cvss:{},risk:{}, cve:{}, bid:{}, cert:{}, xref:{}, description: {}:{}. FULL REPORT AT {}'\
                    .format(self.name, self.family, self.host, self.port, self.cvss, self.risk, self.cve, self.bid, self.cert, self.xref, self.tags, self.description, link)
        logger.debug(message)
        super().__init__(self.vuln_id, type, severity, message)
                    
    
trig_ddl = DDL("CREATE TRIGGER vuln_delete AFTER DELETE ON vulns FOR EACH ROW BEGIN DELETE FROM incidents where incident_id = old.vuln_id; END")
tbl = Vulnerability.__table__
listen(tbl, 'after_create', trig_ddl)     

class Discovery(Incident):
    __tablename__ = 'discovery'
    __mapper_args__ = {'polymorphic_identity': 'discovery'}
    
    discovery_id = Column(GUID, ForeignKey('incidents.incident_id', ondelete='CASCADE'), primary_key=True)
    addr_id = Column('addr_id',  ForeignKey('addrs.addr_id', ondelete='CASCADE'))
    
    cause = Column('cause', String(VARCHAR_LENGTH))
    
    nb_name = Column('nbname', String(VARCHAR_LENGTH))
    nb_info = Column('nb_info', UnicodeText(TEXT_LENGTH))
    dns = Column('dns', String(VARCHAR_LENGTH))
    os = Column('os', String(VARCHAR_LENGTH))
    smb_info = Column('smb_info', UnicodeText(TEXT_LENGTH))
    
    addr = relationship("Address", backref = "discoveries")
    
    def __init__(self, ip, interface, cause):
        self.discovery_id = uuid.uuid1()
        self.cause = cause
        
        self._nmap_info(ip)
        self._dns(ip)
            
        type = incident_by_cause()[cause][0]
        severity = incident_by_cause()[cause][1]

        Incident.__init__(self, self.discovery_id, type, severity, description = 'Discovery of ip: {}, interface: {}, nb_name: {}, os: {}, dns: {}, nbinfo: {}, smbinfo: {}'\
                                                        .format(ip, interface, self.nb_name, self.os, self.dns, self.nb_info, self.smb_info))
        
    def _dns(self,  ip):
        try:
            if DO_DNS is True:
                self.dns = utils.dnsresolve(ip)
            else:
                self.dns = 'Disabled'
        except:
            self.dns = 'None'
            report.dmzwatch_error(  
                                    'DNS error',
                                    Severity.low, 
                                    *sys.exc_info()
                                  )
                
    def _nmap_info(self, ip): 
        try:
            if DO_NBSCAN is True:
                netbios = utils.nbscan('127.0.0.1')
                #netbios = utils.nbscan(ip)
                self.nb_info = netbios['info']
                self.nb_name = netbios['nbname']
            else:
                self.nb_info = 'Disabled'
                self.nb_name = 'Disabled'
        except:
            self.nb_info = 'None'
            self.nb_name = 'None'
            report.dmzwatch_error(  
                                    'Nmap nbscan error',
                                    Severity.low, 
                                    *sys.exc_info()
                                  )
          
        try:
            if DO_SMBSCAN is True:
                smb = utils.smbscan('127.0.0.1')
                #smb = utils.smbscan(ip)
                self.smb_info = smb['info']
                self.os = smb['smbos']
            else:
                self.smb_info = 'Disabled'
                self.os = 'Disabled'
        except:
            self.smb_info = 'None'
            self.nb_name = 'None'
            report.dmzwatch_error(  
                        'Nmap smbscan error',
                        Severity.low, 
                        *sys.exc_info()
                      )
        
 
trig_ddl = DDL("CREATE TRIGGER discovery_delete AFTER DELETE ON discovery FOR EACH ROW BEGIN DELETE FROM incidents where incident_id = old.discovery_id; END")
tbl = Discovery.__table__
listen(tbl, 'after_create', trig_ddl) 

class Address(Base):
    __tablename__ = 'addrs'
    uid = Column('addr_id', Integer, primary_key=True)
    
    ip = Column('ip', String(VARCHAR_LENGTH), nullable=False)
    first_seen = Column('fseen', DateTime, nullable=False)
    last_seen = Column('lseen', DateTime, nullable=False)

    status = Column('status', String(VARCHAR_LENGTH)) 
    interface = Column('interface', String(VARCHAR_LENGTH))
    network = Column('network', String(VARCHAR_LENGTH))
    
    netadmin_id = Column('netadmin_id', Integer)
    
    def __init__(self, ip, interface):
        logger.debug('-----------------------------------')
        logger.debug('New host with ip %s on interface %s', ip, interface) 
        self.ip = ip
        self.interface = interface
        self.first_seen = datetime.datetime.now()
        self.last_seen = self.first_seen
        self.status = HostStatus.active
        
        self._netadmin()
            
        if self.network == Network.unknown:
            cause = Cause.network
        else:
            cause = Cause.new
            
        self.add_discovery(cause)
        self.add_task(cause)
        
    def _netadmin(self):
        try:
            network = netadmin.get_network(self.ip)
        except:
            self.netadmin_id = -1
            self.network = Network.unknown
            
            report.dmzwatch_error(  
                        'Netadmin connection error',
                        Severity.medium, 
                        *sys.exc_info()
                      )
        else:
            self.netadmin_id = network['id']
            self.network = network['network']
        
    def _periodic(self, timeout = DISCOVERY_TIMEOUT ):
        if self.status != HostStatus.active:
            return
        discoveries = self.discoveries
        logger.debug('Checking periodic discovery of %s', self.ip)
        last = max(discoveries, key = lambda disc:disc.date).date
        logger.debug('Latest discovery at: %s', last)
        current = datetime.datetime.now()
        
        if current - last > DISCOVERY_TIMEOUT:
                logger.debug('Adding perioding discovery')
                self._netadmin()
                self.add_discovery(Cause.periodic)
                self.add_task(Cause.periodic)
    
    def seen(self):
        logger.debug('-----------------------------------')
        if self.status == HostStatus.vanished:
            logger.debug('Host with ip %s on interface %s has returned', self.ip, self.interface)

            self.status = HostStatus.returned
            self.add_discovery(Cause.returned)
            self.add_task(Cause.returned)
        else:
            logger.debug('Host with ip %s on interface %s has seen', self.ip, self.interface)
            
            self.status = HostStatus.active
        
        self.last_seen = datetime.datetime.now()
        
        self._periodic()
 

    def aging(self, soft = SOFT_TIMEOUT, hard = HARD_TIMEOUT):
        
        current = datetime.datetime.now()
        diff = current - self.last_seen
        
        if diff > hard and self.status != HostStatus.vanished:
            self.status = HostStatus.vanished 
            logger.debug('-----------------------------------')
            logger.debug('Host with ip %s on interface %s has vanished', self.ip, self.interface)
        elif diff > soft and self.status != HostStatus.inactive and self.status != HostStatus.vanished:
            logger.debug('-----------------------------------')
            logger.debug('Host with ip %s on interface %s become inactive', self.ip, self.interface)
            self.status = HostStatus.inactive
            
    
    def add_task(self, cause):
        logger.debug('Adding standard task to %s wih cause: %s', self.ip, cause )
        
        
        if self.network in task_by_network().keys():
            self.tasks.append(Task(cause, task_by_network()[self.network]))
    
    def add_discovery(self, cause):
        logger.debug('Adding discovery of %s with cause: %s', self.ip, cause)
        self.discoveries.append(Discovery(self.ip, self.interface, cause)) 
        
    
class Task(Base):
    
    __tablename__ = 'tasks'
    
    uid = Column('task_id', Integer, primary_key=True)
    addr_id = Column('addr_id',  ForeignKey('addrs.addr_id', ondelete='CASCADE'))
    start = Column('start', DateTime)
    end = Column('end', DateTime)
    scan_profile = Column('scans', String(VARCHAR_LENGTH))
    cause = Column('cause', String(VARCHAR_LENGTH))
    status = Column('status', String(VARCHAR_LENGTH))
    addr = relationship("Address", backref = 'tasks')
    
    def __init__(self, cause, type, *scans ): #NOT SAFE 
        if type == TaskType.custom:
            assert len(scans) > 0, 'you must provide scan profile if youve choosen type custom '
            self.scan_profile = ','.join(scans)
        else:
            self.scan_profile = scans_by_task()[type]
        logger.debug('-----------------------------------')
        logger.debug('Adding task of %s with scans %s', type, self.scan_profile)
        self.status = ExecStatus.new
        self.cause = cause
    
    def do(self):
        logger.debug('-----------------------------------')
        logger.debug('Doing task for %s', self.addr.ip)
        
        if self.status == ExecStatus.new:
            logger.debug('This task is new')
            for scan in self.scan_profile.split(','):
                self.scans.append(Scan(scan.strip()))
                self.status = ExecStatus.running  
        
        if self.status == ExecStatus.running:
            logger.debug('This task is running now')
            self.status = ExecStatus.complete
            current_scans = 0
            for scan in self.scans:
                if scan.status == ExecStatus.running or scan.status == ExecStatus.new:
                    if current_scans < SIM_SCANS:
                        scan.do()
                        if scan.status == ExecStatus.running:
                            current_scans = current_scans + 1
                    if scan.status == ExecStatus.running or scan.status == ExecStatus.new:
                            self.status = ExecStatus.running
                            
            if self.status == ExecStatus.complete:
                logger.debug('This task for %s complete',  self.addr.ip)
        
            
class Scan(Base):
    __tablename__ = 'scans'
    
    uid = Column('scan_id', Integer, primary_key=True)
    task_id = Column('task_id',  ForeignKey('tasks.task_id', ondelete='CASCADE'))

    start = Column('start', DateTime)
    end = Column('end', DateTime)
    
    type = Column('type', String(VARCHAR_LENGTH))
    status = Column('status', String(VARCHAR_LENGTH))
    
    xml = Column('report_xml', UnicodeText(TEXT_LENGTH))      
    html = Column('report_html', String(VARCHAR_LENGTH))    
    pdf = Column('report_pdf', String(VARCHAR_LENGTH))
    
    info = Column('info', String(VARCHAR_LENGTH))
    extra = Column('extra', PickleType)
    task = relationship("Task", backref = "scans")
    
    
    def __init__(self, type):
        logger.debug('Addded scanning of %s', type)
        self.type = type
        self.status = ExecStatus.new
        
    def do(self):
        try:
            scan_options = scansbytype()[self.type]
        except:
            self.status = ExecStatus.error
            self.info = 'Cannot find options for scan {}.'.format(self.type)
            report.dmzwatch_error(  self.info,
                                    Severity.medium, 
                                    *sys.exc_info()
                                  )
            return
            
        try:
            handler = getattr(Scan, '_' + scan_options[0])
        except:
            self.status = ExecStatus.error
            self.info = 'Cannot find handler for scan {}.'.format(scan_options[0])
            report.dmzwatch_error(  self.info,
                                    Severity.medium, 
                                    *sys.exc_info()
                                  )
            return
            
        arguments = scan_options[1]
        try:
            handler(self, self.task.addr.ip, **arguments)
        except:
            self.status = ExecStatus.error
            self.info = 'Error within scanning {}.'.format(scan_options[0])
            report.dmzwatch_error(  self.info,
                                    Severity.medium, 
                                    *sys.exc_info()
                                  )
    
    
    def _nmap(self, ip, args):
        
        if self.status == ExecStatus.new:
            logger.debug('%s scanning of %s running', self.type, self.task.addr.ip)
            self.start = datetime.datetime.now()
            self.status = ExecStatus.running
            
            try:
                result = nm.nmap_scan('127.0.0.1', args)
            except:
                self.status = ExecStatus.error
                self.end = datetime.datetime.now()
                self.info = 'Nmap scanning error for ip: {}'.format(ip)
                report.dmzwatch_error(  self.info,
                                    Severity.medium, 
                                    *sys.exc_info()
                                  )
                return
               
            self.status = ExecStatus.complete
            self.end = datetime.datetime.now()
            
            self.xml = result['xml']
            
            filename = '{},{},{},{}'.format(self.task.addr.ip, self.type, self.status, self.end.isoformat('_') )
            path = SCAN_REPORT_DIR + filename
            logger.debug(os.getcwd())
            try:
                with open(path + '.xml', 'w') as f:
                    f.write(result['xml'])
                utils.xml_to_html(path + '.xml', path + '.html')
            except:
                self.info = 'Error during nmap report xml/html formatting for {}'.format(ip)
                self.xml = ExecStatus.error
                self.html = ExecStatus.error
                report.dmzwatch_error(  
                                    self.info,
                                    Severity.medium, 
                                    *sys.exc_info()
                                  )
            else:
               
                self.html = SCAN_REPORT_LINK + '/' +  filename + '.html'
                self.pdf = 'NOT IMPLEMENTED'

            for res in result['services']:
                res.update({'link' : self.html})
                self.services.append(Service(**res))
   
        
    def _openvas(self, ip, default_config, default_portlist):
        logger.debug('handle openva scanning for %s', ip)
        if self.status == ExecStatus.new:
            logger.debug('scan is new')
            try:
                a = ovas.openvas_tasks_number()
            except:
                self.start = datetime.datetime.now()
                self.end = datetime.datetime.now()
                self.status = ExecStatus.error
                self.info = 'Cannot get number of tasks!'
                report.dmzwatch_error(  
                    self.info,
                    Severity.medium, 
                    *sys.exc_info()
                  )
                return
            else:
                if a > MAX_OPENVAS_TASKS:
                    logger.debug('MAXIMUM OPENVAS TASKS REACHED')
                    return
             
             
            try:
                self.start = datetime.datetime.now()
                logger.debug('scan is new, adding task to openvas')
                self.extra = ovas.openvas_add_default(ip, default_config, default_portlist)
            except:
                self.end = datetime.datetime.now()
                self.status = ExecStatus.error
                self.info = 'Cannot add openvas task for ip: {}, config: {}, portlist: {}, scan_id: {}'.format(ip, default_config, default_portlist, self.uid)
                report.dmzwatch_error(  
                    self.info,
                    Severity.medium, 
                    *sys.exc_info()
                  )
            else:
                
                self.status = ExecStatus.running
                
        elif self.status == ExecStatus.running:
            logger.debug('scan is running')
            
            try:
                logger.debug('scan if running now, checking status')
                status = ovas.openvas_check_task(self.extra['task'])
                logger.debug('openvas task status is %s', status)
            except:
                self.end = datetime.datetime.now()
                self.status = ExecStatus.error
                self.info = 'Cannot check status of task {}'.format(self.extra)
                report.dmzwatch_error(  
                    self.info,
                    Severity.medium, 
                    *sys.exc_info()
                  )
                return
            
            if self.status == status:
                return
                
                
            logger.debug('Task seems to be complete')
            self.status = status
            self.end = datetime.datetime.now()
            try:
                logger.debug('Getting reports')
                reports = ovas.openvas_get_reports(self.extra['report'])
            except:
                self.info = 'Cannot retrieve reports for ovas task: {}'.format(self.extra)
                report.dmzwatch_error(  
                                self.info,
                                Severity.medium, 
                                *sys.exc_info()
                            )
                     
                return
            
            self.xml = reports['xml'].decode('utf-8')
            self.html = self._save_report('html', reports['html'])
            self.pdf = self._save_report('pdf', reports['html'])
            
            try:
                resvulns = ovas.parse_xml(reports['xml'])
            except:
                self.info = 'Invalid xml report for openvas scanning'
                self.status = ExecStatus.error
                report.dmzwatch_error(  
                                self.info,
                                Severity.medium, 
                                *sys.exc_info()
                            )
            else:
                for vuln in resvulns:
                    self.vulns.append(Vulnerability(link = _get_link(self.html), **vuln))
                
                            
    def _save_report(self, report_type, report_data):
        filename = '{},{},{},{}.{}'.format(self.task.addr.ip, self.type, self.status, self.end.isoformat('_'), report_type)
        path = SCAN_REPORT_DIR + filename
        with open(path, 'bw') as f:
            f.write(report_data)
        
        return filename
        
def _get_link(filename):
        return SCAN_REPORT_LINK + '/' + filename
        
@contextmanager
def session_scope():
    """Provide a transactional scope around a series of operations."""
    session = Session()
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        logger.error('---Something has gone wrong. DMZWATCH DB transaction rollback----')
        raise
    finally:
        session.close()

        
def transactional(fn):
    """Add transactional semantics to a method."""

    def transact(self, *args):
        session = Session()
        try:
            fn(self, session, *args)
            session.commit()
        except:
            session.rollback()
            raise
    transact.__name__ = fn.__name__
    return transact
    
Base.metadata.create_all(engine)
    
